program program_psr
use madx_ptc_module
use pointer_lattice
implicit none

    interface
       subroutine build_PSR(PSR,MIS)
         use madx_ptc_module
         use pointer_lattice
         implicit none
         type(layout), target :: PSR
         logical(lp) mis
       end subroutine build_PSR
    end interface

type(layout), pointer:: psr
type(damap) id,m
type(real_8) ray(6)
real(dp) closed_orbit(6)
type(internal_state) state
type(normalform) normal
integer :: no1=0,np1=0 ,ind(2)=0
type(taylor) z_i_z_j,average_floquet,average_xp
logical(lp) mis

call ptc_ini

psr=>m_u%start

Write(6,*) " small misalignments and rotations in lattice ? input t or f "
read(5,*) mis

call build_PSR(PSR,mis)

!!!  Phase space of Taylor maps will be (x,px,y,py)
state=only_4d0


do while(no1<=1) 
 write(6,*) " input an order > 1 (not too big first)! "
 read(5,*) no1
enddo

np1=0    !!!  No system parameters in the map

call init(state,no1,np1) ! PTC and FPP are properly initialized

call alloc(id,m)  
call alloc(ray)
call alloc(normal)
call alloc(z_i_z_j,average_floquet,average_xp)

closed_orbit=0.d0 ! initial guess for closed orbit

call find_orbit(psr,closed_orbit,1,state,1.d-5)

write(6,'(a16,4(1x,g12.5))') " closed orbit = ",closed_orbit(1:4)

id=1;   ray=closed_orbit+id;   ! ray= closed orbit + identity map

call track(psr,ray,1,state) ! One turn map is computed via the ray

!!! The ray is really the result of a truncated power series algebra
!!! The map we are interested to compute is around the closed orbit: it is really part of a diffential algebra
m=ray   !   The ray is "officially" turned into a "damap"  (DA in Berz's honour)

normal = m  ! The map is normalised   

write(6,*) " Give  i,j for <x_i x_j> "
read(5,*) ind(1),ind(2) 


z_i_z_j=id%v(ind(1))*id%v(ind(2))


!!! This averages an arbitrary function over the orbit
!!! Internally this routine suffers in clarity from the absence of a Complex TPSA  
!!! Input = function z_i_z_j and transformation normal%A_T   -> Normal form = normal%A_T^-1 o Map o normal%A_T
call AVERAGE(z_i_z_j,normal%A_T,average_floquet,average_xp,use_J=.true.)


!!! Output = average_floquet -> average as a function of J_i= (x_i^2 + p_i^2)/2
!!! J_1 is in position 1 and J_2 in position 3
Write(6,*) " average_floquet from output of AVERAGE"
call print(average_floquet,6)

!!! Output = average_xp -> average as a function of x,px,y,py in the ORIGINAL VARIABLES

Write(6,*) " average_xp from output of AVERAGE"
call print(average_xp,6)

!!!! Let us do a test: construct average_xp using average_floquet and normal%A_T^-1 
!!!! First create a special map

m=0

m%v(1)= (id%v(1)**2+id%v(2)**2)/2.d0   !   J_1
m%v(3)= (id%v(3)**2+id%v(4)**2)/2.d0   !   J_2

m=m*normal%A_T**(-1)   ! going in the original variables I_1= J_1 *normal%A_T**(-1)    SEE SECTION ON INVARIANTS

average_xp=0

average_xp=average_floquet*m

Write(6,*) " average_xp reconstructed using output of AVERAGE"
call print(average_xp,6)

write(6,*) "   "
write(6,*) " hit return to terminate program "
write(6,*) "   "
pause 
call kill(id,m)  
call kill(ray)
call kill(normal)
call kill(z_i_z_j,average_floquet,average_xp)

end program program_psr


!=================================================================
subroutine  build_PSR(PSR,mis)
use madx_ptc_module
use pointer_lattice
implicit none

type(layout), target :: PSR

real(dp) :: ang, brho, kd, kf, Larc,SG(6),cut
type(fibre) :: b, d1, d2, qd, qf
type(layout) :: cell
logical(lp) mis
!-----------------------------------

call make_states(.false.)
exact_model = .true.
default = default + nocavity  
call update_states
madlength = .false.

ang = (twopi * 36.d0 / 360.d0)
Larc = 2.54948d0
brho = 1.2d0 * (Larc / ang)
call set_mad(brho = brho, method = 2, step = 10)
madkind2 = drift_kick_drift

kf =  2.72d0 / brho
kd = -1.92d0 / brho

d1 = drift("D1", 2.28646d0)
d2 = drift("D2", 0.45d0)
qf = quadrupole("QF", 0.5d0, kf)
qd = quadrupole("QD", 0.5d0, kd)
b  = rbend("B", Larc, ang)
cell = d1 + qd + d2 + b + d2 + qf + d1

PSR = 10 * cell
PSR = .ring.PSR

call survey(PSR)

if(mis) then
 sig=1.d-5
 cut=4.d0
 call MESS_UP_ALIGNMENT(PSR,SIG,cut)
endif
end subroutine build_PSR
